#!/usr/bin/env python3
import argparse
import csv
import json
import os
from typing import Dict, Tuple


def load_template(template_path: str) -> Dict[str, str]:
    """Load caption -> category mapping from template CSV.

    Expects columns: caption, category (and optional count).
    Uses exact caption string match.
    """
    mapping: Dict[str, str] = {}
    with open(template_path, 'r', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        for row in reader:
            caption = (row.get('caption') or '').strip()
            category = (row.get('category') or '').strip()
            if caption and category:
                mapping[caption] = category
    return mapping


def build_image_path(image_root: str, image_filename: str) -> str:
    # image_filename is expected like: image_00030334.jpg
    image_filename = (image_filename or '').strip()
    return os.path.join(image_root, image_filename)


def process_metadata(metadata_path: str,
                     template_map: Dict[str, str],
                     image_root: str,
                     out_csv: str,
                     include_unlabeled: bool = False,
                     unlabeled_label: str = '') -> Tuple[int, int, int]:
    """Generate labeled CSV from metadata using caption->category mapping.

    Writes columns: image, text, label.
    Returns tuple: (total, labeled, skipped)
    """
    with open(metadata_path, 'r', encoding='utf-8') as f:
        meta = json.load(f)

    total = 0
    labeled = 0
    skipped = 0

    # Ensure parent directory exists
    os.makedirs(os.path.dirname(out_csv), exist_ok=True)

    with open(out_csv, 'w', encoding='utf-8', newline='') as f_out:
        writer = csv.writer(f_out)
        writer.writerow(['image', 'text', 'label'])

        for item in meta:
            total += 1
            caption = (item.get('caption') or '').strip()
            image_filename = item.get('image_filename') or ''
            image_path = build_image_path(image_root, image_filename)

            category = template_map.get(caption)
            if category:
                writer.writerow([image_path, caption, category])
                labeled += 1
            else:
                if include_unlabeled:
                    writer.writerow([image_path, caption, unlabeled_label])
                else:
                    skipped += 1

    return total, labeled, skipped


def main():
    parser = argparse.ArgumentParser(description='Apply caption labels to metadata and generate labeled CSV.')
    parser.add_argument('--metadata', required=True, help='Path to metadata.json')
    parser.add_argument('--template', required=True, help='Path to caption->category template CSV')
    parser.add_argument('--image_root', required=True, help='Root directory of images to prefix file_name')
    parser.add_argument('--out_csv', required=True, help='Output CSV path')
    parser.add_argument('--include_unlabeled', action='store_true', help='Include rows without labels with empty label or provided --unlabeled_label')
    parser.add_argument('--unlabeled_label', default='', help='Label to use for unlabeled rows when --include_unlabeled is set')
    args = parser.parse_args()

    template_map = load_template(args.template)
    total, labeled, skipped = process_metadata(
        metadata_path=args.metadata,
        template_map=template_map,
        image_root=args.image_root,
        out_csv=args.out_csv,
        include_unlabeled=args.include_unlabeled,
        unlabeled_label=args.unlabeled_label,
    )

    print(f'Template entries: {len(template_map)}')
    print(f'Total metadata items: {total}')
    print(f'Labeled rows written: {labeled}')
    print(f'Skipped (unlabeled): {skipped}')


if __name__ == '__main__':
    main()